import { TreeDataNode } from 'antd';

import { DatasetCollectionTypeEnum, TrainingModeEnum } from './enum';

export interface TreeData extends TreeDataNode {
  parentKey?: string;
  parent?: TreeData;
  type?: string;
  fileId?: string;
  children?: TreeData[];
  name?: string;
  id?: string;
  updateTime?: string;
  createTime?: string;
  count?: number;
  forbid?: boolean;
  agent_model?: string;
  app_id?: string;
  avatar?: string;
  create_time?: Date;
  dataset_type?: Array<string>;
  default_permission?: string;
  group_id?: string;
  inherit_permission?: boolean;
  intro?: string;
  metadata?: CollectionMetadataType;
  openapiUrl?: string;
  parent_id?: string | null;
  status?: string;
  team_id?: string;
  tmb_id?: string;
  update_time?: Date;
  vector_model?: string;
  reviewStatus?: string;
  collectionId?: string | undefined;
  appId?: string;
  groupId?: string;
  _id?: string;
  userId?: string;
  operator?: string;
  experiodDays?: string | null;
  experiodTimes?: string | null;
  groupName?: string;
}

export enum DatasetEvent {
  UPDATECOLLECTION,
  UPDATEDATASETTREE,
  UPDATECHILDREN,
  DATASETTREECLICKED,
}

export enum CollectionEvent {
  UPDATECOLLECTION,
  UPDATECOLLECTIONDATAS,
}

export type DatasetFileSchema = {
  _id: string;
  length: number;
  chunkSize: number;
  uploadDate: Date;
  filename: string;
  contentType: string;
  metadata: {
    teamId: string;
    tmbId: string;
    encoding?: string;
  };
};

export type CollectionMetadataType = {
  webPageSelector?: string;
  relatedImgId?: string; // The id of the associated image collections
  hashFileBuffers?: string;
  publishTime?: Date;
  expiredTime?: string;

  [key: string]: any;
};

export type AddCollectionMetadataType = {
  collection_id: string;
  metadata: CollectionMetadataType;
};

export type DelCollectionMetadataType = {
  collection_id: string;
  metadata: string[];
};

// collection
export type DatasetCollectionSchemaType = {
  sourceName: string;
  sourceId?: string;
  file?: DatasetFileSchema;
  // permission: DatasetPermission;
  _id: string;
  teamId: string;
  tmbId: string;
  datasetId: string;
  parentId?: string;
  name: string;
  type: DatasetCollectionTypeEnum;
  createTime: Date;
  updateTime: Date;
  forbid?: boolean;

  trainingType: TrainingModeEnum;
  chunkSize: number;
  chunkSplitter?: string;
  qaPrompt?: string;
  ocrParse?: boolean;

  tags?: string[];

  fileId?: string; // local file id
  rawLink?: string; // link url
  externalFileId?: string; // external file id

  rawTextLength?: number;
  hashRawText?: string;
  externalFileUrl?: string; // external import url
  metadata?: CollectionMetadataType;

  indexes?: string[];
  reviewStatus?: string;
};

export type DatasetCollectionsListItemType = {
  _id: string;
  parentId?: DatasetCollectionSchemaType['parentId'];
  tmbId: DatasetCollectionSchemaType['tmbId'];
  name: DatasetCollectionSchemaType['name'];
  type: DatasetCollectionSchemaType['type'];
  createTime: DatasetCollectionSchemaType['createTime'];
  updateTime: DatasetCollectionSchemaType['updateTime'];
  forbid?: DatasetCollectionSchemaType['forbid'];
  trainingType?: DatasetCollectionSchemaType['trainingType'];
  tags?: string[];

  externalFileId?: string;

  fileId?: string;
  rawLink?: string;

  dataAmount: number;
  trainingAmount: number;
};

export type DatasetDataListItemType = {
  _id: string;
  datasetId: string;
  collectionId: string;
  q: string; // embedding content
  a: string; // bonus content
  chunkIndex?: number;
  dataId?: string;
  // indexes: DatasetDataSchemaType['indexes'];
};

export enum SearchScoreTypeEnum {
  embedding = 'embedding',
  fullText = 'fullText',
  reRank = 'reRank',
  rrf = 'rrf',
}

export enum DatasetSearchModeEnum {
  embedding = 'embedding',
  fullTextRecall = 'fullTextRecall',
  mixedRecall = 'mixedRecall',
  knowledgeGraph = 'knowledgeGraph',
}
export interface GetDatasetReq {
  app_id: string;
  isPage: boolean;
  dataset_type: string;
  parentId?: string;
}

export interface GetDatasetResList {
  _id: string;
  parentId: null;
  teamId: string;
  tmbId: string;
  type: string;
  status: string;
  avatar: string;
  name: string;
  vectorModel: string;
  agentModel: string;
  intro: string;
  inheritPermission: true;
  updateTime: string;
  __v: number;
  id: string;
}

export interface GetDatasetRes {
  total: number;
  list: GetDatasetResList[];
}
